{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import re\n",
    "from IPython.core.debugger import set_trace\n",
    "from pprint import pprint\n",
    "import unicodedata\n",
    "from transformers import BertModel, BasicTokenizer, BertTokenizerFast\n",
    "import copy\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torch.optim as optim\n",
    "import glob\n",
    "import time\n",
    "from common.utils import Preprocessor\n",
    "from tplinker_plus import (HandshakingTaggingScheme,\n",
    "                          DataMaker4Bert, \n",
    "                          DataMaker4BiLSTM, \n",
    "                          TPLinkerPlusBert, \n",
    "                          TPLinkerPlusBiLSTM,\n",
    "                          MetricsCalculator)\n",
    "import wandb\n",
    "import config\n",
    "from glove import Glove\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = config.eval_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(config[\"device_num\"])\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_home = config[\"data_home\"]\n",
    "experiment_name = config[\"exp_name\"]\n",
    "test_data_path = os.path.join(data_home, experiment_name, config[\"test_data\"])\n",
    "batch_size = config[\"hyper_parameters\"][\"batch_size\"]\n",
    "rel2id_path = os.path.join(data_home, experiment_name, config[\"rel2id\"])\n",
    "ent2id_path = os.path.join(data_home, experiment_name, config[\"ent2id\"])\n",
    "save_res_dir = os.path.join(config[\"save_res_dir\"], experiment_name)\n",
    "max_seq_len = config[\"hyper_parameters\"][\"max_seq_len\"]\n",
    "# for reproductivity\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_path_dict = {}\n",
    "for file_path in glob.glob(test_data_path):\n",
    "    file_name = re.search(\"(.*?)\\.json\", file_path.split(\"/\")[-1]).group(1)\n",
    "    test_data_path_dict[file_name] = file_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data_dict = {}\n",
    "for file_name, path in test_data_path_dict.items():\n",
    "    test_data_dict[file_name] = json.load(open(path, \"r\", encoding = \"utf-8\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "if config[\"encoder\"] == \"BERT\":\n",
    "    tokenizer = BertTokenizerFast.from_pretrained(config[\"bert_path\"], add_special_tokens = False, do_lower_case = False)\n",
    "    tokenize = tokenizer.tokenize\n",
    "    get_tok2char_span_map = lambda text: tokenizer.encode_plus(text, return_offsets_mapping = True, add_special_tokens = False)[\"offset_mapping\"]\n",
    "elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "    tokenize = lambda text: text.split(\" \")\n",
    "    def get_tok2char_span_map(text):\n",
    "        tokens = text.split(\" \")\n",
    "        tok2char_span = []\n",
    "        char_num = 0\n",
    "        for tok in tokens:\n",
    "            tok2char_span.append((char_num, char_num + len(tok)))\n",
    "            char_num += len(tok) + 1 # +1: whitespace\n",
    "        return tok2char_span"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocessor = Preprocessor(tokenize_func = tokenize, \n",
    "                            get_tok2char_span_map_func = get_tok2char_span_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Calculate the max token number: 100%|██████████| 824/824 [00:00<00:00, 9828.82it/s]\n"
     ]
    }
   ],
   "source": [
    "all_data = []\n",
    "for data in list(test_data_dict.values()):\n",
    "    all_data.extend(data)\n",
    "    \n",
    "max_tok_num = 0\n",
    "for sample in tqdm(all_data, desc = \"Calculate the max token number\"):\n",
    "    tokens = tokenize(sample[\"text\"])\n",
    "    max_tok_num = max(len(tokens), max_tok_num)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max_tok_num: 138, less than or equal to max_test_seq_len: 512, no need to split!\n"
     ]
    }
   ],
   "source": [
    "split_test_data = False\n",
    "if max_tok_num > max_seq_len:\n",
    "    split_test_data = True\n",
    "    print(\"max_tok_num: {}, lagger than max_test_seq_len: {}, test data will be split!\".format(max_tok_num, max_seq_len))\n",
    "else:\n",
    "    print(\"max_tok_num: {}, less than or equal to max_test_seq_len: {}, no need to split!\".format(max_tok_num, max_seq_len))\n",
    "max_seq_len = min(max_tok_num, max_seq_len) \n",
    "\n",
    "if config[\"hyper_parameters\"][\"force_split\"]:\n",
    "    split_test_data = True\n",
    "    print(\"force to split the test dataset!\")    \n",
    "\n",
    "ori_test_data_dict = copy.deepcopy(test_data_dict)\n",
    "if split_test_data:\n",
    "    test_data_dict = {}\n",
    "    for file_name, data in ori_test_data_dict.items():\n",
    "        test_data_dict[file_name] = preprocessor.split_into_short_samples(data, \n",
    "                                                                          max_seq_len, \n",
    "                                                                          sliding_len = config[\"hyper_parameters\"][\"sliding_len\"], \n",
    "                                                                          encoder = config[\"encoder\"], \n",
    "                                                                          data_type = \"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Decoder(Tagger)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "rel2id = json.load(open(rel2id_path, \"r\", encoding = \"utf-8\"))\n",
    "ent2id = json.load(open(ent2id_path, \"r\", encoding = \"utf-8\"))\n",
    "handshaking_tagger = HandshakingTaggingScheme(rel2id, max_seq_len, ent2id)\n",
    "tag_size = handshaking_tagger.get_tag_size()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "if config[\"encoder\"] == \"BERT\":\n",
    "    tokenizer = BertTokenizerFast.from_pretrained(config[\"bert_path\"], add_special_tokens = False, do_lower_case = False)\n",
    "    data_maker = DataMaker4Bert(tokenizer, handshaking_tagger)\n",
    "    get_tok2char_span_map = lambda text: tokenizer.encode_plus(text, return_offsets_mapping = True, add_special_tokens = False)[\"offset_mapping\"]\n",
    "\n",
    "elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "    token2idx_path = os.path.join(*config[\"token2idx_path\"])\n",
    "    token2idx = json.load(open(token2idx_path, \"r\", encoding = \"utf-8\"))\n",
    "    idx2token = {idx:tok for tok, idx in token2idx.items()}\n",
    "    def text2indices(text, max_seq_len):\n",
    "        input_ids = []\n",
    "        tokens = text.split(\" \")\n",
    "        for tok in tokens:\n",
    "            if tok not in token2idx:\n",
    "                input_ids.append(token2idx['<UNK>'])\n",
    "            else:\n",
    "                input_ids.append(token2idx[tok])\n",
    "        if len(input_ids) < max_seq_len:\n",
    "            input_ids.extend([token2idx['<PAD>']] * (max_seq_len - len(input_ids)))\n",
    "        input_ids = torch.tensor(input_ids[:max_seq_len])\n",
    "        return input_ids\n",
    "    def get_tok2char_span_map(text):\n",
    "        tokens = text.split(\" \")\n",
    "        tok2char_span = []\n",
    "        char_num = 0\n",
    "        for tok in tokens:\n",
    "            tok2char_span.append((char_num, char_num + len(tok)))\n",
    "            char_num += len(tok) + 1 # +1: whitespace\n",
    "        return tok2char_span\n",
    "    data_maker = DataMaker4BiLSTM(text2indices, get_tok2char_span_map, handshaking_tagger)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        return self.data[index]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "if config[\"encoder\"] == \"BERT\":\n",
    "    encoder = BertModel.from_pretrained(config[\"bert_path\"])\n",
    "    hidden_size = encoder.config.hidden_size\n",
    "    rel_extractor = TPLinkerPlusBert(encoder, \n",
    "                                     tag_size,\n",
    "                                     config[\"hyper_parameters\"][\"shaking_type\"],\n",
    "                                     config[\"hyper_parameters\"][\"inner_enc_type\"],\n",
    "                                    )\n",
    "    \n",
    "elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "    glove = Glove()\n",
    "    glove = glove.load(config[\"pretrained_word_embedding_path\"])\n",
    "    \n",
    "    # prepare embedding matrix\n",
    "    word_embedding_init_matrix = np.random.normal(-1, 1, size=(len(token2idx), config[\"word_embedding_dim\"]))\n",
    "    count_in = 0\n",
    "\n",
    "    # 在预训练词向量中的用该预训练向量\n",
    "    # 不在预训练集里的用随机向量\n",
    "    for ind, tok in tqdm(idx2token.items(), desc=\"Embedding matrix initializing...\"):\n",
    "        if tok in glove.dictionary:\n",
    "            count_in += 1\n",
    "            word_embedding_init_matrix[ind] = glove.word_vectors[glove.dictionary[tok]]\n",
    "\n",
    "    print(\"{:.4f} tokens are in the pretrain word embedding matrix\".format(count_in / len(idx2token))) # 命中预训练词向量的比例\n",
    "    word_embedding_init_matrix = torch.FloatTensor(word_embedding_init_matrix)\n",
    "    \n",
    "    rel_extractor = TPLinkerPlusBiLSTM(word_embedding_init_matrix, \n",
    "                                       config[\"emb_dropout\"], \n",
    "                                       config[\"enc_hidden_size\"], \n",
    "                                       config[\"dec_hidden_size\"],\n",
    "                                       config[\"rnn_dropout\"],\n",
    "                                       tag_size, \n",
    "                                       config[\"hyper_parameters\"][\"shaking_type\"],\n",
    "                                       config[\"hyper_parameters\"][\"inner_enc_type\"],\n",
    "                                      )\n",
    "\n",
    "rel_extractor = rel_extractor.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Merics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = MetricsCalculator(handshaking_tagger)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get model state paths\n",
    "model_state_dir = config[\"model_state_dict_dir\"]\n",
    "target_run_ids = set(config[\"run_ids\"])\n",
    "run_id2model_state_paths = {}\n",
    "for root, dirs, files in os.walk(model_state_dir):\n",
    "    for file_name in files:\n",
    "        run_id = root[-8:]\n",
    "        if re.match(\".*model_state.*\\.pt\", file_name) and run_id in target_run_ids:\n",
    "            if run_id not in run_id2model_state_paths:\n",
    "                run_id2model_state_paths[run_id] = []\n",
    "            model_state_path = os.path.join(root, file_name)\n",
    "            run_id2model_state_paths[run_id].append(model_state_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_last_k_paths(path_list, k):\n",
    "    path_list = sorted(path_list, key = lambda x: int(re.search(\"(\\d+)\", x.split(\"/\")[-1]).group(1)))\n",
    "#     pprint(path_list)\n",
    "    return path_list[-k:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# only last k models\n",
    "k = config[\"last_k_model\"]\n",
    "for run_id, path_list in run_id2model_state_paths.items():\n",
    "    run_id2model_state_paths[run_id] = get_last_k_paths(path_list, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_duplicates(rel_list, ent_list):\n",
    "    rel_memory_set = set()\n",
    "    filtered_rel_list = []\n",
    "    \n",
    "    for rel in rel_list:\n",
    "        rel_memory = \"{}\\u2E80{}\\u2E80{}\\u2E80{}\\u2E80{}\".format(rel[\"subj_tok_span\"][0], \n",
    "                                                                 rel[\"subj_tok_span\"][1], \n",
    "                                                                 rel[\"predicate\"], \n",
    "                                                                 rel[\"obj_tok_span\"][0], \n",
    "                                                                 rel[\"obj_tok_span\"][1])\n",
    "        if rel_memory not in rel_memory_set:\n",
    "            filtered_rel_list.append(rel)\n",
    "            rel_memory_set.add(rel_memory)\n",
    "            \n",
    "    ent_memory_set = set()\n",
    "    filtered_ent_list = []\n",
    "    for ent in ent_list:\n",
    "        ent_memory = \"{}\\u2E80{}\\u2E80{}\".format(ent[\"tok_span\"][0], \n",
    "                                                 ent[\"tok_span\"][1], \n",
    "                                                 ent[\"type\"])\n",
    "        if ent_memory not in ent_memory_set:\n",
    "            filtered_ent_list.append(ent)\n",
    "            ent_memory_set.add(ent_memory)\n",
    "            \n",
    "    return filtered_rel_list, filtered_ent_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(test_data, ori_test_data):\n",
    "    '''\n",
    "    test_data: if split, it would be samples with subtext\n",
    "    ori_test_data: the original data has not been split, used to get original text here\n",
    "    '''\n",
    "    indexed_test_data = data_maker.get_indexed_data(test_data, max_seq_len, data_type = \"test\") # fill up to max_seq_len\n",
    "    test_dataloader = DataLoader(MyDataset(indexed_test_data), \n",
    "                              batch_size = batch_size, \n",
    "                              shuffle = False, \n",
    "                              num_workers = 6,\n",
    "                              drop_last = False,\n",
    "                              collate_fn = lambda data_batch: data_maker.generate_batch(data_batch, data_type = \"test\"),\n",
    "                             )\n",
    "    \n",
    "    pred_sample_list = []\n",
    "    for batch_test_data in tqdm(test_dataloader, desc = \"Predicting\"):\n",
    "        if config[\"encoder\"] == \"BERT\":\n",
    "            sample_list, batch_input_ids, \\\n",
    "            batch_attention_mask, batch_token_type_ids, \\\n",
    "            tok2char_span_list, _ = batch_test_data\n",
    "\n",
    "            batch_input_ids, \\\n",
    "            batch_attention_mask, \\\n",
    "            batch_token_type_ids = (batch_input_ids.to(device), \n",
    "                                      batch_attention_mask.to(device), \n",
    "                                      batch_token_type_ids.to(device),\n",
    "                                     )\n",
    "\n",
    "        elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "            sample_list, batch_input_ids, \\\n",
    "            tok2char_span_list, _ = batch_test_data\n",
    "\n",
    "            batch_input_ids = batch_input_ids.to(device)\n",
    "            \n",
    "        with torch.no_grad():\n",
    "            if config[\"encoder\"] == \"BERT\":\n",
    "                batch_pred_shaking_tag, _ = rel_extractor(batch_input_ids, \n",
    "                                                          batch_attention_mask, \n",
    "                                                          batch_token_type_ids, \n",
    "                                                         )\n",
    "            elif config[\"encoder\"] in {\"BiLSTM\", }:\n",
    "                batch_pred_shaking_tag, _ = rel_extractor(batch_input_ids)\n",
    "\n",
    "        batch_pred_shaking_tag = (batch_pred_shaking_tag > 0.).long()\n",
    "        \n",
    "        for ind in range(len(sample_list)):\n",
    "            gold_sample = sample_list[ind]\n",
    "            text = gold_sample[\"text\"]\n",
    "            text_id = gold_sample[\"id\"]\n",
    "            tok2char_span = tok2char_span_list[ind]\n",
    "            pred_shaking_tag = batch_pred_shaking_tag[ind]\n",
    "            tok_offset, char_offset = 0, 0\n",
    "            if split_test_data:\n",
    "                tok_offset, char_offset = gold_sample[\"tok_offset\"], gold_sample[\"char_offset\"]\n",
    "            rel_list, ent_list = handshaking_tagger.decode_rel(text,\n",
    "                                                     pred_shaking_tag,\n",
    "                                                     tok2char_span, \n",
    "                                                     tok_offset = tok_offset, char_offset = char_offset)\n",
    "            pred_sample_list.append({\n",
    "                \"text\": text,\n",
    "                \"id\": text_id,\n",
    "                \"relation_list\": rel_list,\n",
    "                \"entity_list\": ent_list,\n",
    "            })\n",
    "            \n",
    "    # merge\n",
    "    text_id2pred_res = {}\n",
    "    for sample in pred_sample_list:\n",
    "        text_id = sample[\"id\"]\n",
    "        if text_id not in text_id2pred_res:\n",
    "            text_id2pred_res[text_id] = {\n",
    "                \"rel_list\": sample[\"relation_list\"],\n",
    "                \"ent_list\": sample[\"entity_list\"],\n",
    "            }\n",
    "        else:\n",
    "            text_id2pred_res[text_id][\"rel_list\"].extend(sample[\"relation_list\"])\n",
    "            text_id2pred_res[text_id][\"ent_list\"].extend(sample[\"entity_list\"])\n",
    "\n",
    "    text_id2text = {sample[\"id\"]:sample[\"text\"] for sample in ori_test_data}\n",
    "    merged_pred_sample_list = []\n",
    "    for text_id, pred_res in text_id2pred_res.items():\n",
    "        filtered_rel_list, filtered_ent_list = filter_duplicates(pred_res[\"rel_list\"], pred_res[\"ent_list\"])\n",
    "        merged_pred_sample_list.append({\n",
    "            \"id\": text_id,\n",
    "            \"text\": text_id2text[text_id],\n",
    "            \"relation_list\": filtered_rel_list,\n",
    "            \"entity_list\": filtered_ent_list,\n",
    "        })\n",
    "        \n",
    "    return merged_pred_sample_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_test_prf(pred_sample_list, gold_test_data, pattern = \"whole_text\"):\n",
    "    text_id2gold_n_pred = {} # text id to gold and pred results\n",
    "    \n",
    "    for sample in gold_test_data:\n",
    "        text_id = sample[\"id\"]\n",
    "        text_id2gold_n_pred[text_id] = {\n",
    "            \"gold_relation_list\": sample[\"relation_list\"],\n",
    "            \"gold_entity_list\": sample[\"entity_list\"],\n",
    "            \"gold_event_list\": sample[\"event_list\"],\n",
    "        }\n",
    "    \n",
    "    for sample in pred_sample_list:\n",
    "        text_id = sample[\"id\"]\n",
    "        text_id2gold_n_pred[text_id][\"pred_relation_list\"] = sample[\"relation_list\"]\n",
    "        text_id2gold_n_pred[text_id][\"pred_entity_list\"] = sample[\"entity_list\"]\n",
    "\n",
    "    correct_num, pred_num, gold_num = 0, 0, 0\n",
    "    ent_correct_num, ent_pred_num, ent_gold_num = 0, 0, 0\n",
    "    ee_cpg_dict = {\n",
    "        \"trigger_iden_cpg\": [0, 0, 0],\n",
    "        \"trigger_class_cpg\": [0, 0, 0],\n",
    "        \"arg_iden_cpg\": [0, 0, 0],\n",
    "        \"arg_class_cpg\": [0, 0, 0],\n",
    "        }\n",
    "    ere_cpg_dict = {\n",
    "            \"rel_cpg\": [0, 0, 0],\n",
    "            \"ent_cpg\": [0, 0, 0],\n",
    "            }\n",
    "    for gold_n_pred in text_id2gold_n_pred.values():\n",
    "        gold_rel_list = gold_n_pred[\"gold_relation_list\"]\n",
    "        pred_rel_list = gold_n_pred[\"pred_relation_list\"] if \"pred_relation_list\" in gold_n_pred else []\n",
    "        gold_ent_list = gold_n_pred[\"gold_entity_list\"]\n",
    "        pred_ent_list = gold_n_pred[\"pred_entity_list\"] if \"pred_entity_list\" in gold_n_pred else []\n",
    "\n",
    "        if pattern == \"event_extraction\":\n",
    "            pred_event_list = handshaking_tagger.trans2ee(pred_rel_list, pred_ent_list) # transform to event list\n",
    "            gold_event_list = gold_n_pred[\"gold_event_list\"] # *\n",
    "            metrics.cal_event_cpg(pred_event_list, gold_event_list, ee_cpg_dict)\n",
    "        else:\n",
    "            metrics.cal_rel_cpg(pred_rel_list, pred_ent_list, gold_rel_list, gold_ent_list, ere_cpg_dict, pattern)\n",
    "       \n",
    "    if pattern == \"event_extraction\":\n",
    "        trigger_iden_prf = metrics.get_prf_scores(ee_cpg_dict[\"trigger_iden_cpg\"][0], \n",
    "                                                  ee_cpg_dict[\"trigger_iden_cpg\"][1], \n",
    "                                                  ee_cpg_dict[\"trigger_iden_cpg\"][2])\n",
    "        trigger_class_prf = metrics.get_prf_scores(ee_cpg_dict[\"trigger_class_cpg\"][0], \n",
    "                                                  ee_cpg_dict[\"trigger_class_cpg\"][1], \n",
    "                                                  ee_cpg_dict[\"trigger_class_cpg\"][2])\n",
    "        arg_iden_prf = metrics.get_prf_scores(ee_cpg_dict[\"arg_iden_cpg\"][0], ee_cpg_dict[\"arg_iden_cpg\"][1], ee_cpg_dict[\"arg_iden_cpg\"][2])\n",
    "        arg_class_prf = metrics.get_prf_scores(ee_cpg_dict[\"arg_class_cpg\"][0], ee_cpg_dict[\"arg_class_cpg\"][1], ee_cpg_dict[\"arg_class_cpg\"][2])\n",
    "        prf_dict = {\n",
    "            \"trigger_iden_prf\": trigger_iden_prf,\n",
    "            \"trigger_class_prf\": trigger_class_prf,\n",
    "            \"arg_iden_prf\": arg_iden_prf,\n",
    "            \"arg_class_prf\": arg_class_prf,\n",
    "        }\n",
    "        return prf_dict\n",
    "    else:\n",
    "        rel_prf = metrics.get_prf_scores(ere_cpg_dict[\"rel_cpg\"][0], ere_cpg_dict[\"rel_cpg\"][1], ere_cpg_dict[\"rel_cpg\"][2])\n",
    "        ent_prf = metrics.get_prf_scores(ere_cpg_dict[\"ent_cpg\"][0], ere_cpg_dict[\"ent_cpg\"][1], ere_cpg_dict[\"ent_cpg\"][2])\n",
    "        prf_dict = {\n",
    "            \"rel_prf\": rel_prf,\n",
    "            \"ent_prf\": ent_prf,\n",
    "        }\n",
    "        return prf_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate indexed train or valid data: 97it [00:00, 576.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "run_id: 1a70p109, model state model_state_dict_3.pt loaded\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Generate indexed train or valid data: 824it [00:00, 2479.01it/s]\n",
      "Predicting: 100%|██████████| 26/26 [00:28<00:00,  1.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'../results/ace05_lu/1a70p109/test_data_res_3.json': 368}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# predict\n",
    "res_dict = {}\n",
    "predict_statistics = {}\n",
    "for file_name, short_data in test_data_dict.items():\n",
    "    ori_test_data = ori_test_data_dict[file_name]\n",
    "    for run_id, model_path_list in run_id2model_state_paths.items():\n",
    "        save_dir4run = os.path.join(save_res_dir, run_id)\n",
    "        if config[\"save_res\"] and not os.path.exists(save_dir4run):\n",
    "            os.makedirs(save_dir4run)\n",
    "            \n",
    "        for model_state_path in model_path_list:\n",
    "            res_num = re.search(\"(\\d+)\", model_state_path.split(\"/\")[-1]).group(1)\n",
    "            save_path = os.path.join(save_dir4run, \"{}_res_{}.json\".format(file_name, res_num))\n",
    "            \n",
    "            if os.path.exists(save_path):\n",
    "                pred_sample_list = [json.loads(line) for line in open(save_path, \"r\", encoding = \"utf-8\")]\n",
    "                print(\"{} already exists, load it directly!\".format(save_path))\n",
    "            else:\n",
    "                # load model state\n",
    "                rel_extractor.load_state_dict(torch.load(model_state_path))\n",
    "                rel_extractor.eval()\n",
    "                print(\"run_id: {}, model state {} loaded\".format(run_id, model_state_path.split(\"/\")[-1]))\n",
    "\n",
    "                # predict\n",
    "                pred_sample_list = predict(short_data, ori_test_data)\n",
    "            \n",
    "            res_dict[save_path] = pred_sample_list\n",
    "            predict_statistics[save_path] = len([s for s in pred_sample_list if len(s[\"relation_list\"]) > 0])\n",
    "pprint(predict_statistics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "check char span: 100%|██████████| 824/824 [00:00<00:00, 209055.56it/s]\n"
     ]
    }
   ],
   "source": [
    "# check\n",
    "for path, res in res_dict.items():\n",
    "    for sample in tqdm(res, desc = \"check char span\"):\n",
    "        text = sample[\"text\"]\n",
    "        for rel in sample[\"relation_list\"]:\n",
    "            assert rel[\"subject\"] == text[rel[\"subj_char_span\"][0]:rel[\"subj_char_span\"][1]]\n",
    "            assert rel[\"object\"] == text[rel[\"obj_char_span\"][0]:rel[\"obj_char_span\"][1]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save \n",
    "if config[\"save_res\"]:\n",
    "    for path, res in res_dict.items():\n",
    "        with open(path, \"w\", encoding = \"utf-8\") as file_out:\n",
    "            for sample in tqdm(res, desc = \"Output\"):\n",
    "                if len(sample[\"relation_list\"]) == 0:\n",
    "                    continue\n",
    "                json_line = json.dumps(sample, ensure_ascii = False)     \n",
    "                file_out.write(\"{}\\n\".format(json_line))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> \u001b[0;32m/root/wycheng/research/tplinker/tplinker_plus/tplinker_plus.py\u001b[0m(262)\u001b[0;36mtrans2ee\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m    260 \u001b[0;31m            \u001b[0;32mif\u001b[0m \u001b[0mtirigger_offset2event\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtrigger_offset_str\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mevent_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m# filter false relations\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    261 \u001b[0;31m                \u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m--> 262 \u001b[0;31m                \u001b[0;32mcontinue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    263 \u001b[0;31m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m    264 \u001b[0;31m            \u001b[0;31m# append arguments\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n",
      "ipdb> event_type\n",
      "'Transport'\n",
      "ipdb> tirigger_offset2event[trigger_offset_str]\n",
      "'Attack'\n",
      "ipdb> tirigger_offset2score[trigger_offset_str]\n",
      "*** NameError: name 'tirigger_offset2score' is not defined\n",
      "ipdb> tirigger_offset2vote[trigger_offset_str]\n",
      "*** NameError: name 'tirigger_offset2vote' is not defined\n",
      "ipdb> trigger_offset2vote\n",
      "{'6,7': {'Attack': 4.1}, '11,12': {'Attack': 4.1, 'Transport': 1}}\n",
      "ipdb> rel_list\n",
      "[{'subject': 'marines', 'object': 'battled', 'subj_tok_span': [4, 6], 'obj_tok_span': [6, 7], 'subj_char_span': [5, 12], 'obj_char_span': [13, 20], 'predicate': 'Attacker_Attack'}, {'subject': 'snipers', 'object': 'battled', 'subj_tok_span': [7, 9], 'obj_tok_span': [6, 7], 'subj_char_span': [21, 28], 'obj_char_span': [13, 20], 'predicate': 'Attacker_Attack'}, {'subject': 'snipers', 'object': 'battled', 'subj_tok_span': [7, 9], 'obj_tok_span': [6, 7], 'subj_char_span': [21, 28], 'obj_char_span': [13, 20], 'predicate': 'Target_Attack'}, {'subject': 'snipers', 'object': 'fought', 'subj_tok_span': [7, 9], 'obj_tok_span': [11, 12], 'subj_char_span': [21, 28], 'obj_char_span': [37, 43], 'predicate': 'Attacker_Attack'}, {'subject': 'they', 'object': 'fought', 'subj_tok_span': [10, 11], 'obj_tok_span': [11, 12], 'subj_char_span': [32, 36], 'obj_char_span': [37, 43], 'predicate': 'Attacker_Attack'}, {'subject': 'capital', 'object': 'fought', 'subj_tok_span': [15, 16], 'obj_tok_span': [11, 12], 'subj_char_span': [60, 67], 'obj_char_span': [37, 43], 'predicate': 'Destination_Transport'}, {'subject': 'capital', 'object': 'fought', 'subj_tok_span': [15, 16], 'obj_tok_span': [11, 12], 'subj_char_span': [60, 67], 'obj_char_span': [37, 43], 'predicate': 'Place_Attack'}]\n",
      "ipdb> ent_list\n",
      "[{'type': 'Argument_Attacker', 'text': 'marines', 'tok_span': [4, 6], 'char_span': [5, 12]}, {'type': 'Argument_Target', 'text': 'marines', 'tok_span': [4, 6], 'char_span': [5, 12]}, {'type': 'Trigger_Attack', 'text': 'battled', 'tok_span': [6, 7], 'char_span': [13, 20]}, {'type': 'Argument_Attacker', 'text': 'snipers', 'tok_span': [7, 9], 'char_span': [21, 28]}, {'type': 'Argument_Target', 'text': 'snipers', 'tok_span': [7, 9], 'char_span': [21, 28]}, {'type': 'Argument_Artifact', 'text': 'they', 'tok_span': [10, 11], 'char_span': [32, 36]}, {'type': 'Argument_Attacker', 'text': 'they', 'tok_span': [10, 11], 'char_span': [32, 36]}, {'type': 'Trigger_Attack', 'text': 'fought', 'tok_span': [11, 12], 'char_span': [37, 43]}, {'type': 'Argument_Destination', 'text': 'capital', 'tok_span': [15, 16], 'char_span': [60, 67]}, {'type': 'Argument_Place', 'text': 'capital', 'tok_span': [15, 16], 'char_span': [60, 67]}]\n",
      "ipdb> q\n"
     ]
    },
    {
     "ename": "BdbQuit",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mBdbQuit\u001b[0m                                   Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-25-86a01941a1df>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      5\u001b[0m         \u001b[0mfile_name\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mre\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msearch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"(.*?)_res_\\d+\\.json\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfile_path\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"/\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      6\u001b[0m         \u001b[0mgold_test_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mori_test_data_dict\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mfile_name\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 7\u001b[0;31m         \u001b[0mprf_dict\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_test_prf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_samples\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgold_test_data\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpattern\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconfig\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"hyper_parameters\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"match_pattern\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      8\u001b[0m         \u001b[0mfilepath2scores\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mfile_path\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprf_dict\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      9\u001b[0m     \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"---------------- Results -----------------------\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-21-1803cd982783>\u001b[0m in \u001b[0;36mget_test_prf\u001b[0;34m(pred_sample_list, gold_test_data, pattern)\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     35\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mpattern\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"event_extraction\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 36\u001b[0;31m             \u001b[0mpred_event_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhandshaking_tagger\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrans2ee\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_rel_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpred_ent_list\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# transform to event list\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     37\u001b[0m             \u001b[0mgold_event_list\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgold_n_pred\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"gold_event_list\"\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;31m# *\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     38\u001b[0m             \u001b[0mmetrics\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcal_event_cpg\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_event_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgold_event_list\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mee_cpg_dict\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/root/wycheng/research/tplinker/tplinker_plus/tplinker_plus.py\u001b[0m in \u001b[0;36mtrans2ee\u001b[0;34m(self, rel_list, ent_list)\u001b[0m\n\u001b[1;32m    260\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mtirigger_offset2event\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtrigger_offset_str\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mevent_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m# filter false relations\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    261\u001b[0m                 \u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m                 \u001b[0;32mcontinue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    263\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    264\u001b[0m             \u001b[0;31m# append arguments\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/root/wycheng/research/tplinker/tplinker_plus/tplinker_plus.py\u001b[0m in \u001b[0;36mtrans2ee\u001b[0;34m(self, rel_list, ent_list)\u001b[0m\n\u001b[1;32m    260\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mtirigger_offset2event\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtrigger_offset_str\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mevent_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;31m# filter false relations\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    261\u001b[0m                 \u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m                 \u001b[0;32mcontinue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    263\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    264\u001b[0m             \u001b[0;31m# append arguments\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.8/bdb.py\u001b[0m in \u001b[0;36mtrace_dispatch\u001b[0;34m(self, frame, event, arg)\u001b[0m\n\u001b[1;32m     86\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0;31m# None\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     87\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'line'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 88\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     89\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mevent\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m'call'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     90\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdispatch_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.8/bdb.py\u001b[0m in \u001b[0;36mdispatch_line\u001b[0;34m(self, frame)\u001b[0m\n\u001b[1;32m    111\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstop_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbreak_here\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    112\u001b[0m             \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0muser_line\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mframe\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 113\u001b[0;31m             \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mquitting\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mBdbQuit\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    114\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_dispatch\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    115\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mBdbQuit\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# score\n",
    "if config[\"score\"]:\n",
    "    filepath2scores = {}\n",
    "    for file_path, pred_samples in res_dict.items():\n",
    "        file_name = re.search(\"(.*?)_res_\\d+\\.json\", file_path.split(\"/\")[-1]).group(1)\n",
    "        gold_test_data = ori_test_data_dict[file_name]\n",
    "        prf_dict = get_test_prf(pred_samples, gold_test_data, pattern = config[\"hyper_parameters\"][\"match_pattern\"])\n",
    "        filepath2scores[file_path] = prf_dict\n",
    "    print(\"---------------- Results -----------------------\")\n",
    "    pprint(filepath2scores)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
